#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
清理目录团队分配数据

用于清理directories表中的团队分配数据，重新开始分配流程。
"""

import os
import sqlite3
import logging
from datetime import datetime

# 配置日志
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(message)s')
logger = logging.getLogger(__name__)

class DirectoryAssignmentCleaner:
    """目录团队分配清理器"""

    def __init__(self, db_path: str = None):
        """初始化清理器"""
        if db_path is None:
            # 默认数据库路径
            current_dir = os.path.dirname(os.path.abspath(__file__))
            project_root = os.path.dirname(os.path.dirname(current_dir))
            db_path = os.path.join(project_root, 'backend', 'data', 'uma_insight.db')

        self.db_path = db_path

    def connect_database(self) -> sqlite3.Connection:
        """连接数据库"""
        try:
            conn = sqlite3.connect(self.db_path)
            return conn
        except Exception as e:
            logger.error(f"连接数据库失败: {e}")
            raise

    def backup_before_clear(self, conn: sqlite3.Connection):
        """清理前备份数据"""
        timestamp = datetime.now().strftime('%Y%m%d_%H%M%S')
        backup_table_name = f"directories_assignment_backup_{timestamp}"

        logger.info(f"创建分配数据备份: {backup_table_name}")
        conn.execute(f"""
            CREATE TABLE {backup_table_name} AS
            SELECT
                id, path, name, metric_count, team_id,
                assignment_confidence, assigned_at, assignment_method
            FROM directories
            WHERE team_id IS NOT NULL
        """)
        conn.commit()

        # 查询备份的记录数
        cursor = conn.cursor()
        cursor.execute(f"SELECT COUNT(*) FROM {backup_table_name}")
        backup_count = cursor.fetchone()[0]
        logger.info(f"备份了 {backup_count} 条分配记录到表 {backup_table_name}")

        return backup_table_name

    def clear_assignments(self, conn: sqlite3.Connection, create_backup: bool = True):
        """清理团队分配数据"""
        logger.info("开始清理目录团队分配数据...")

        if create_backup:
            self.backup_before_clear(conn)

        # 清理directories表中的分配相关字段
        cursor = conn.cursor()
        cursor.execute("""
            UPDATE directories
            SET team_id = NULL,
                assignment_confidence = 0,
                assigned_at = NULL,
                assignment_method = NULL
            """)

        affected_rows = cursor.rowcount
        conn.commit()

        logger.info(f"成功清理了 {affected_rows} 个目录的团队分配数据")
        return affected_rows

    def verify_clear(self, conn: sqlite3.Connection):
        """验证清理结果"""
        logger.info("验证清理结果...")

        cursor = conn.cursor()
        cursor.execute("SELECT COUNT(*) FROM directories WHERE team_id IS NOT NULL")
        remaining_assignments = cursor.fetchone()[0]

        cursor.execute("SELECT COUNT(*) FROM directories")
        total_directories = cursor.fetchone()[0]

        if remaining_assignments == 0:
            logger.info("✓ 清理验证通过：所有目录的团队分配已清空")
            logger.info(f"总目录数: {total_directories}")
            return True
        else:
            logger.warning(f"✗ 清理验证失败：仍有 {remaining_assignments} 个目录存在团队分配")
            return False

    def list_backup_tables(self, conn: sqlite3.Connection):
        """列出备份表"""
        logger.info("检查现有的备份表...")

        cursor = conn.cursor()
        cursor.execute("""
            SELECT name, sql FROM sqlite_master
            WHERE type='table' AND name LIKE 'directories_assignment_backup_%'
            ORDER BY name DESC
        """)

        backup_tables = cursor.fetchall()
        if backup_tables:
            logger.info(f"找到 {len(backup_tables)} 个备份表:")
            for table_name, sql in backup_tables:
                cursor.execute(f"SELECT COUNT(*) FROM {table_name}")
                count = cursor.fetchone()[0]
                logger.info(f"  {table_name}: {count} 条记录")
        else:
            logger.info("未找到备份表")

        return backup_tables

    def cleanup_old_backups(self, conn: sqlite3.Connection, keep_count: int = 3):
        """清理旧的备份表"""
        logger.info(f"清理旧的备份表，保留最近 {keep_count} 个...")

        cursor = conn.cursor()
        cursor.execute("""
            SELECT name FROM sqlite_master
            WHERE type='table' AND name LIKE 'directories_assignment_backup_%'
            ORDER BY name DESC
        """)

        backup_tables = [row[0] for row in cursor.fetchall()]
        old_backups = backup_tables[keep_count:]

        if old_backups:
            logger.info(f"将删除 {len(old_backups)} 个旧备份表:")
            for table_name in old_backups:
                cursor.execute(f"DROP TABLE {table_name}")
                logger.info(f"  已删除: {table_name}")

            conn.commit()
            logger.info("旧备份表清理完成")
        else:
            logger.info("没有需要清理的旧备份表")

        return len(old_backups)

    def run_clear(self, create_backup: bool = True, cleanup_backups: bool = True):
        """执行完整的清理流程"""
        logger.info("开始目录团队分配数据清理流程")

        try:
            conn = self.connect_database()

            # 列出现有备份表
            self.list_backup_tables(conn)

            # 清理分配数据
            affected_rows = self.clear_assignments(conn, create_backup)

            # 验证清理结果
            success = self.verify_clear(conn)

            if success:
                # 清理旧备份
                if cleanup_backups:
                    self.cleanup_old_backups(conn)

                logger.info("目录团队分配数据清理完成！")
                return True, affected_rows
            else:
                logger.error("目录团队分配数据清理验证失败")
                return False, affected_rows

        except Exception as e:
            logger.error(f"清理过程出错: {e}")
            return False, 0
        finally:
            conn.close()


def main():
    """主函数"""
    import argparse

    parser = argparse.ArgumentParser(description='清理目录团队分配数据工具')
    parser.add_argument('--db-path', help='数据库文件路径')
    parser.add_argument('--no-backup', action='store_true',
                       help='不创建备份（谨慎使用）')
    parser.add_argument('--keep-backups', type=int, default=3,
                       help='保留备份表数量 (默认: 3)')
    parser.add_argument('--list-only', action='store_true',
                       help='仅列出备份表，不执行清理')
    parser.add_argument('--cleanup-backups-only', action='store_true',
                       help='仅清理旧备份表')

    args = parser.parse_args()

    try:
        cleaner = DirectoryAssignmentCleaner(args.db_path)
        conn = cleaner.connect_database()

        if args.list_only:
            cleaner.list_backup_tables(conn)

        elif args.cleanup_backups_only:
            deleted_count = cleaner.cleanup_old_backups(conn, args.keep_backups)
            logger.info(f"删除了 {deleted_count} 个旧备份表")

        else:
            # 执行清理
            create_backup = not args.no_backup
            cleanup_backups = True

            if not create_backup:
                logger.warning("⚠️  警告：将不创建备份直接清理数据")

            success, affected_rows = cleaner.run_clear(create_backup, cleanup_backups)

            if success:
                logger.info("清理任务完成")
                return 0
            else:
                logger.error("清理任务失败")
                return 1

        conn.close()

    except Exception as e:
        logger.error(f"执行过程出错: {e}")
        return 1

    return 0


if __name__ == "__main__":
    exit(main())